'''
    网络 U2Net
    数据集 MKDataset
    优化器：Adam
    损失函数：对比次数6+1次， 二分类损失 BCELoss()
    保存模型：
'''
import test_20250318.u2net as u2net
import data_loader
from torch.utils.data import DataLoader
import torch
from tqdm import tqdm

# 损失函数
bec_loss = torch.nn.BCELoss()


def loss_fun(d0, d1, d2, d3, d4, d5, d6, labels):
    loss0 = bec_loss(d0, labels)
    loss1 = bec_loss(d1, labels)
    loss2 = bec_loss(d2, labels)
    loss3 = bec_loss(d3, labels)
    loss4 = bec_loss(d4, labels)
    loss5 = bec_loss(d5, labels)
    loss6 = bec_loss(d6, labels)
    return loss0 + loss6 + loss5 + loss4 + loss2 + loss3 + loss1


if __name__ == '__main__':
    # 网络
    net = u2net.U2NET().cuda()
    # 数据集
    train_loader = DataLoader(dataset=data_loader.MKDataset(path="data"), batch_size=4, shuffle=True)
    test_loader = DataLoader(dataset=data_loader.MKDataset(path="data", isTrain=False), batch_size=4, shuffle=True)
    # 优化器
    opt = torch.optim.Adam(net.parameters())

    # 开始训练
    EPOCHS = 100
    for epoch in range(EPOCHS):
        net.train()
        loss_sum = 0.
        count = 0
        for img, labels in tqdm(train_loader):
            img, labels = img.cuda(), labels.cuda()
            d0, d1, d2, d3, d4, d5, d6 = net(img)
            loss = loss_fun(d0, d1, d2, d3, d4, d5, d6, labels)
            loss_sum += loss.item()
            count += 1
            opt.zero_grad()
            loss.backward()
            opt.step()

        avg_loss = loss_sum / count
        print(f"epoch: {epoch} , avg_loss : {avg_loss}")

